{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e017d77b",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/AI-Hypercomputer/maxtext/blob/main/src/MaxText/examples/demo_decoding.ipynb)\n",
    "        \n",
    "# Qwen3-0.6B Decoding Demo"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc85cefe-8f29-47db-a8f3-4e8fbb354eb5",
   "metadata": {},
   "source": [
    "## Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55e3ce9e-8968-4d68-ba2b-b36c616b52a9",
   "metadata": {},
   "source": [
    "### Change Runtime Type\n",
    "\n",
    "**Instructions:**\n",
    "1.  Navigate to the menu at the top of the screen.\n",
    "2.  Click on **Runtime**.\n",
    "3.  Select **Change runtime type** from the dropdown menu.\n",
    "4.  Select **v5e-1 TPU** as the **Hardware accelerator**.\n",
    "5. Click on **Save**."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf5e0f3f-5833-4260-a31d-b156249d67ab",
   "metadata": {},
   "source": [
    "### Get Your Hugging Face Token\n",
    "\n",
    "To access model checkpoint from the Hugging Face Hub, you need to authenticate with a personal access token.\n",
    "\n",
    "**Follow these steps to get your token:**\n",
    "\n",
    "1.  **Navigate to the Access Tokens page** in your Hugging Face account settings. You can go there directly by visiting this URL:\n",
    "    *   [https://huggingface.co/settings/tokens](https://huggingface.co/settings/tokens)\n",
    "\n",
    "2.  **Create a new token** by clicking the **\"+ Create new token\"** button.\n",
    "\n",
    "3.  **Give your token a name** and assign it a **`read` role**. The `read` role is sufficient for downloading models.\n",
    "\n",
    "4.  **Copy the generated token**. You will need to paste it in the next step.\n",
    "\n",
    "**Follow these steps to store your token:**\n",
    "\n",
    "1. On the left sidebar of your Colab window, click the key icon (the Secrets tab).\n",
    "\n",
    "2. Click **\"+ Add new secret\"**.\n",
    "\n",
    "3. Set the Name as **HF_TOKEN**.\n",
    "\n",
    "4. Paste your token into the Value field.\n",
    "\n",
    "5. Ensure the Notebook access toggle is turned On."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a6deec5-b64a-4bc6-86c4-c24696c66f17",
   "metadata": {},
   "source": [
    "## Installation: MaxText & Other Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b2d4a66-99de-404c-aac3-18b1af4af78e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install uv, a fast Python package installer\n",
    "!pip install uv\n",
    "\n",
    "# Install MaxText and dependencies\n",
    "!uv pip install maxtext --resolution=lowest\n",
    "!python3 -m MaxText.install_maxtext_extra_deps\n",
    "\n",
    "# Use nest_asyncio to allow nested event loops in notebooks\n",
    "!uv pip install nest_asyncio\n",
    "\n",
    "# Install the PyTorch library\n",
    "!uv pip install torch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a07fd61-35b7-4aa9-93cd-49ef89fb550d",
   "metadata": {},
   "source": [
    "### Restart Session\n",
    "To apply certain changes, you need to restart the session.\n",
    "\n",
    "**Instructions:**\n",
    "1.  Navigate to the menu at the top of the screen.\n",
    "2.  Click on **Runtime**.\n",
    "3.  Select **Restart session** from the dropdown menu.\n",
    "\n",
    "You will be asked to confirm the action in a pop-up dialog. Click on **Yes**."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f1ebdb1-dcf4-417b-9c29-3461e06aa9cf",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8e986cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "import datetime\n",
    "import jax\n",
    "import os\n",
    "import nest_asyncio\n",
    "import numpy as np\n",
    "\n",
    "import MaxText as mt\n",
    "from MaxText import common_types\n",
    "from MaxText import inference_utils\n",
    "from MaxText import maxtext_utils\n",
    "from MaxText import max_logging\n",
    "from MaxText import pyconfig\n",
    "from MaxText.input_pipeline import _input_pipeline_utils\n",
    "from MaxText.utils.ckpt_conversion import to_maxtext\n",
    "\n",
    "from google.colab import userdata\n",
    "from huggingface_hub import login\n",
    "\n",
    "MAXTEXT_PKG_DIR = os.path.dirname(mt.__file__)\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4f53124",
   "metadata": {},
   "source": [
    "## Sanity Test: Checking for Available TPU Devices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a545acd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "jax.distributed.initialize()  # distributed.initialize should only be called once.\n",
    "jax.devices()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be0113d9-0cb6-45aa-9fa4-7e543db7645e",
   "metadata": {},
   "source": [
    "## Model Configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b80080fa-473b-4683-b0c9-765af43efd49",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = \"qwen3-0.6b\"\n",
    "PROMPT = \"I love to\"\n",
    "RUN_NAME = datetime.datetime.now().strftime(\"%Y-%m-%d-%H-%M-%S\")\n",
    "MODEL_CHECKPOINT_PATH = f\"/tmp/checkpoints/{MODEL_NAME}/{RUN_NAME}/unscanned\"\n",
    "\n",
    "HF_TOKEN = userdata.get(\"HF_TOKEN\")\n",
    "login(token=HF_TOKEN)\n",
    "max_logging.log(\"Authenticated with Hugging Face successfully!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03ff53b8-b931-4190-bcac-d6ca885cbbc8",
   "metadata": {},
   "source": [
    "## Download Model Checkpoint From Hugging Face"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fd3578c-5763-410e-8b61-72d7415628bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "argv = [\n",
    "    \"\",  # This is a placeholder, it's not actually used by the script's logic\n",
    "    f\"{MAXTEXT_PKG_DIR}/configs/base.yml\",\n",
    "    f\"model_name={MODEL_NAME}\",\n",
    "    f\"base_output_directory={MODEL_CHECKPOINT_PATH}\",\n",
    "    f\"hf_access_token={HF_TOKEN}\",\n",
    "    \"use_multimodal=false\",\n",
    "    \"scan_layers=false\",\n",
    "]\n",
    "\n",
    "to_maxtext.main(argv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94a9fd37-95e5-4075-837e-de0f1666d55f",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_logging.log(f\"Model checkpoint can be found at: {MODEL_CHECKPOINT_PATH}/0/items\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cf4bbe4-6485-4ce7-8aef-cf3df3810e52",
   "metadata": {},
   "source": [
    "## Initialize Configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32f44079-87ae-4ed4-a008-c0dbb8aaf8c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "config = pyconfig.initialize(\n",
    "    [\"\", f\"{MAXTEXT_PKG_DIR}/configs/base.yml\"],\n",
    "    per_device_batch_size=1.0,\n",
    "    run_name=\"test\",\n",
    "    max_target_length=4,\n",
    "    max_prefill_predict_length=4,\n",
    "    tokenizer_path=f\"{MAXTEXT_PKG_DIR}/assets/qwen3-tokenizer\",\n",
    "    load_parameters_path=f\"{MODEL_CHECKPOINT_PATH}/0/items\",\n",
    "    model_name=MODEL_NAME,\n",
    "    async_checkpointing=False,\n",
    "    prompt=PROMPT,\n",
    "    scan_layers=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a637f43b-6fcc-4305-af5b-d9d30d464bb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_logging.log(\"Decode configurations initialized.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd502094-1694-410a-91e9-25bbb8dfb33a",
   "metadata": {},
   "source": [
    "## Initialize Decode State"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2d2de93",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = mt.from_config(config)\n",
    "mesh = model.mesh\n",
    "init_rng = jax.random.PRNGKey(config.init_weights_seed)\n",
    "state, _ = maxtext_utils.setup_decode_state(model, config, init_rng, mesh, None)\n",
    "max_logging.log(\"Decode state initialized.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed4b59a7",
   "metadata": {},
   "source": [
    "## Get Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35584129-3c45-45ad-b2a2-a56f98d27f06",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = _input_pipeline_utils.get_tokenizer(\n",
    "    f\"{MAXTEXT_PKG_DIR}/assets/qwen3-tokenizer\",\n",
    "    \"huggingface\",\n",
    "    add_bos=True,\n",
    "    add_eos=False,\n",
    ")\n",
    "max_logging.log(\"Tokenizer loaded succuessfully.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32a252ae",
   "metadata": {},
   "source": [
    "## Prepare Inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2d2d0c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_ids = tokenizer.encode(config.prompt)\n",
    "\n",
    "# Pad input_ids to max_target_length\n",
    "padded_ids = np.zeros(config.max_target_length, dtype=np.int32)\n",
    "padded_ids[: len(input_ids)] = input_ids\n",
    "ids = np.asarray(padded_ids, dtype=np.int32)\n",
    "\n",
    "s = (config.global_batch_size_to_train_on, config.max_target_length)\n",
    "decoder_segment_ids = np.zeros(s) + common_types.DECODING_ACTIVE_SEQUENCE_INDICATOR\n",
    "decoder_positions = np.stack(\n",
    "    [np.arange(config.max_target_length, dtype=np.int32) for _ in range(config.global_batch_size_to_train_on)]\n",
    ")\n",
    "\n",
    "ids = np.stack([ids for _ in range(config.global_batch_size_to_train_on)])\n",
    "max_logging.log(\n",
    "    f\"input_ids={input_ids}, \\n\\nids={ids}, \\n\\ndecoder_segment_ids = {decoder_segment_ids}, \\n\\ndecoder_positions= {decoder_positions}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "647018c1",
   "metadata": {},
   "source": [
    "## Run Forward Pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7436751b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "full_train_logits = model.apply(\n",
    "    state.params,\n",
    "    ids,\n",
    "    decoder_positions,\n",
    "    decoder_segment_ids,\n",
    "    enable_dropout=False,\n",
    "    rngs={\"aqt\": init_rng},\n",
    ")\n",
    "full_train_logits = jax.experimental.multihost_utils.process_allgather(full_train_logits)\n",
    "max_logging.log(f\"{full_train_logits[0, 0, :]=}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5640ab55",
   "metadata": {},
   "source": [
    "## Generate Text with Greedy Decoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb06c0c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "selected_logits = jax.lax.dynamic_slice(\n",
    "    full_train_logits, (0, 0, full_train_logits.shape[2] - 2, 0), (1, 1, 1, full_train_logits.shape[3])\n",
    ")\n",
    "\n",
    "# Consider the greedily sampled token\n",
    "init_rng, new_rng = jax.random.split(init_rng)\n",
    "first_generated_token = inference_utils.sampling(\n",
    "    selected_logits,\n",
    "    new_rng,\n",
    "    config.decode_sampling_strategy,  # \"greedy\"\n",
    ")\n",
    "output = tokenizer.decode([first_generated_token.item()])\n",
    "max_logging.log(f\"Next predicted token is `{output}` for the input prompt: `{config.prompt}`.\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
